package cn.youfule.rtsp2ws;


import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author
 * @version
 */
@Slf4j
public class WsSessionManager {
    /**
     * 存放所有在线的客户端
     */
    private static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();
    /**
     * 保存视频源进行播放的所有客户端session
     */
    private static ConcurrentHashMap<String, Set<String>> SOURCE_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 session
     *
     * @param sessionKey
     * @param sourceKey
     */
    public static void add(String sessionKey, String sourceKey,WebSocketSession session) {
        // 添加 session
        SESSION_POOL.put(sessionKey, session);
        Set<String> set = SOURCE_POOL.get(sourceKey);
        if(set == null){
            set = new HashSet<>();
        }
        set.add(sessionKey);
        SOURCE_POOL.put(sourceKey,set);
    }

    /**
     * 删除 session,会返回删除的 session
     *
     * @param sessionKey
     * @param sourceKey
     * @return
     */
    public static WebSocketSession remove(String sessionKey, String sourceKey) {
        Set<String> set = SOURCE_POOL.get(sourceKey);
        if(set != null){
            set.remove(sessionKey);
            SOURCE_POOL.put(sourceKey,set);
        }
        // 删除 session
        return SESSION_POOL.remove(sessionKey);
    }

    /**
     * 删除并同步关闭连接
     *
     * @param sessionKey
     * @param sourceKey
     */
    public static void removeAndClose(String sessionKey, String sourceKey) {
        WebSocketSession session = remove(sessionKey,sourceKey);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                // todo: 关闭出现异常处理
                e.printStackTrace();
            }
        }
    }

    /**
     * 获得 session
     *
     * @param sessionKey
     * @return
     */
    public static WebSocketSession get(String sessionKey) {
        // 获得 session
        return SESSION_POOL.get(sessionKey);
    }

    /**
     * 获得 session
     *
     * @param sourceKey
     * @return
     */
    public static Set<WebSocketSession> getSessions(String sourceKey) {
        Set<WebSocketSession> res = new HashSet<>();
        Set<String> keys = SOURCE_POOL.get(sourceKey);
        if(keys == null){
            return res;
        }
        // 获得 session
        for (String key:keys){
            res.add(SESSION_POOL.get(key));
        }
        return res;
    }
    /**
     * 获得 session
     *
     * @param sourceKey
     * @return
     */
    public static Set<String> getKeys(String sourceKey) {
        return SOURCE_POOL.get(sourceKey);
    }
}